# 基于matplotlib 使用 yfinance ohlc数据 实现中性杠杆网格回测

import yfinance as yf
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from matplotlib.dates import DateFormatter

from mjr import MACDKDJRSI

class MJRStrategy:
    def __init__(
        self,
        ticker,
        initial_capital=100000,
        grid_range=0.2,
        grid_num=10,
        leverage=3,
        fee=0.0005,
        start="2020-01-01",
        end="2025-02-23",
        period="1y",
        intervals="1d",
    ):
        self.ticker = ticker
        self.initial_capital = initial_capital
        self.grid_range = grid_range  # 价格波动范围
        self.grid_num = grid_num  # 网格层数
        self.leverage = leverage  # 单层杠杆倍数
        self.fee = fee  # 交易费率
        self.start = start
        self.end = end
        self.period = period
        self.intervals = intervals

        self.data = self._download_data()
        # self._generate_data()
        self._run_backtest()

    def _download_data(self):
        mjr = MACDKDJRSI(period=self.period,interval=self.intervals)
        return mjr.df
        # df = yf.download(self.ticker, start=self.start, end=self.end)
        # return df[["Open", "High", "Low", "Close"]]

    def _generate_data(self):
        """动态生成网格体系"""
        base_price = self.data["Close"].iloc
        self.grid_levels = np.linspace(
            base_price * (1 - self.grid_range),
            base_price * (1 + self.grid_range),
            self.grid_num,
        )
        self.grid_spacing = self.grid_levels - self.grid_levels

    def _run_backtest(self):
        """中性策略核心逻辑"""
        capital = self.initial_capital
        position = {"long": 0, "short": 0}  # 多空对冲仓位
        equity = [capital]
        trade_log = []

        for close_price in self.data["Close"]:
            # 计算价格偏离度
            deviation = (close_price - self.grid_levels.mean()) / self.grid_spacing

            # 杠杆平衡公式
            target_long = max(0, deviation) * self.leverage
            target_short = max(0, -deviation) * self.leverage

            # 执行调仓
            delta_long = target_long - position["long"]
            delta_short = target_short - position["short"]

            # 计算交易成本
            fee = abs(delta_long + delta_short) * close_price * self.fee
            capital -= fee

            # 更新仓位和资金
            position = {"long": target_long, "short": target_short}
            equity.append(
                capital + (position["long"] - position["short"]) * close_price
            )

            trade_log.append(
                {
                    "price": close_price,
                    "long": target_long,
                    "short": target_short,
                    "fee": fee,
                }
            )

        self.equity = equity
        self.trade_log = pd.DataFrame(trade_log)

    def visualize(self):
        """可视化模块"""
        plt.figure(figsize=(14, 8))

        # 价格与网格体系
        ax1 = plt.subplot2grid((4, 1), (0, 0), rowspan=3)
        ax1.plot(self.data["Close"], label="Price", linewidth=1)
        for level in self.grid_levels:
            ax1.axhline(level, color="gray", linestyle="--", alpha=0.5)
        ax1.set_title(f"Neutral Grid Strategy - {self.ticker}")

        # 多空仓位
        ax2 = plt.subplot2grid((4, 1), (3, 0), sharex=ax1)
        ax2.plot(self.trade_log["long"], label="Long Position", color="green")
        ax2.plot(self.trade_log["short"], label="Short Position", color="red")
        ax2.set_ylabel("Leverage")

        # 格式化显示
        ax1.xaxis.set_major_formatter(DateFormatter("%Y-%m"))
        plt.gcf().autofmt_xdate()
        ax1.legend()
        ax2.legend()
        plt.tight_layout()

        # 资金曲线
        plt.figure(figsize=(14, 4))
        plt.plot(self.equity, label="Equity Curve")
        plt.title(
            f"Total Return: {((self.equity[-1] / self.initial_capital - 1) * 100):.2f}%"
        )
        plt.xlabel("Trading Days")
        plt.show()


# 策略实例化与运行
strategy = MJRStrategy(
    ticker="BTC-USD",
    initial_capital=100000,
    grid_range=0.3,
    grid_num=20,
    leverage=5,
    start="2024-02-15",
)
strategy.visualize()
